import numpy as np

def calculate_iou(box1, box2):
    """
    计算两个边界框的IoU (Intersection over Union)。
    box: [x1, y1, x2, y2]
    """
    # 提取坐标
    x1_b1, y1_b1, x2_b1, y2_b1 = box1
    x1_b2, y1_b2, x2_b2, y2_b2 = box2

    # 计算交集区域的坐标
    x1_inter = max(x1_b1, x1_b2)
    y1_inter = max(y1_b1, y1_b2)
    x2_inter = min(x2_b1, x2_b2)
    y2_inter = min(y2_b1, y2_b2)

    # 计算交集区域的面积
    inter_width = max(0, x2_inter - x1_inter)
    inter_height = max(0, y2_inter - y1_inter)
    inter_area = inter_width * inter_height

    # 计算两个边界框的面积
    area_b1 = (x2_b1 - x1_b1) * (y2_b1 - y1_b1)
    area_b2 = (x2_b2 - x1_b2) * (y2_b2 - y1_b2)

    # 计算并集面积
    union_area = area_b1 + area_b2 - inter_area

    if union_area == 0:
        return 0.0
    
    iou = inter_area / union_area
    return iou

def calculate_ciou(box1, box2):
    """
    计算两个边界框的CIoU (Complete-IoU)。
    box: [x1, y1, x2, y2]
    """
    iou = calculate_iou(box1, box2)

    # 如果IoU已经接近1，说明框已经完美匹配，直接返回1
    # 避免后续的nan问题
    if iou >= 1.0: # 考虑到浮点数精度，也可以用 iou > 0.9999
        return 1.0

    # 1. 计算中心点距离惩罚项 (DIoU部分)
    x1_b1, y1_b1, x2_b1, y2_b1 = box1
    x1_b2, y1_b2, x2_b2, y2_b2 = box2

    # 计算中心点坐标
    center_x1 = (x1_b1 + x2_b1) / 2
    center_y1 = (y1_b1 + y2_b1) / 2
    center_x2 = (x1_b2 + x2_b2) / 2
    center_y2 = (y1_b2 + y2_b2) / 2

    # 计算中心点欧氏距离的平方
    distance_squared = (center_x1 - center_x2)**2 + (center_y1 - center_y2)**2

    # 计算最小外接矩形C的坐标
    x1_c = min(x1_b1, x1_b2)
    y1_c = min(y1_b1, y1_b2)
    x2_c = max(x2_b1, x2_b2)
    y2_c = max(y2_b1, y2_b2)

    # 计算最小外接矩形C的对角线长度的平方
    c_diag_squared = (x2_c - x1_c)**2 + (y2_c - y1_c)**2

    # 避免除以零
    if c_diag_squared == 0: # 如果外接框退化为点，则距离惩罚项设为0
        distance_penalty = 0.0
    else:
        distance_penalty = distance_squared / c_diag_squared

    # 2. 计算宽高比一致性惩罚项
    # 预测框宽高
    w1 = x2_b1 - x1_b1
    h1 = y2_b1 - y1_b1
    # 真实框宽高
    w2 = x2_b2 - x1_b2
    h2 = y2_b2 - y1_b2

    # 避免除以零，如果宽度或高度为0，则设置一个非常小的值
    epsilon = 1e-6
    w1 = max(epsilon, w1)
    h1 = max(epsilon, h1)
    w2 = max(epsilon, w2)
    h2 = max(epsilon, h2)

    # 计算v
    v = (4 / (np.pi**2)) * (np.arctan(w2 / h2) - np.arctan(w1 / h1))**2

    # 计算 alpha
    # 这里的条件很重要，避免 (1 - iou) + v == 0 导致的 nan
    # 如果 iou 已经为1，v 为0，则 alpha 的分母为0。
    # 实际上，当 IoU = 1 时，所有惩罚项都应该为 0，CIoU = 1。
    # 我们已经在函数开头进行了处理，但为了严谨性，这里再加一个对分母的检查。
    denominator_alpha = (1 - iou) + v
    if abs(denominator_alpha) < epsilon: # 如果分母非常接近0
        alpha = 0.0 # 理论上此时 v 也应为0，所以 alpha v = 0
    else:
        alpha = v / denominator_alpha
    
    # 组合各项
    ciou = iou - distance_penalty - alpha * v
    return ciou

if __name__ == "__main__":
    # 定义一些示例边界框
    # box格式: [x1, y1, x2, y2]

    # 场景1：两个框完全重合 (期待IoU=1, CIoU=1)
    box_a = [0, 0, 10, 10]
    box_b = [0, 0, 10, 10]
    iou_ab = calculate_iou(box_a, box_b)
    ciou_ab = calculate_ciou(box_a, box_b)
    print(f"场景1：框A {box_a}, 框B {box_b} (完全重合)")
    print(f"IoU: {iou_ab:.4f}, CIoU: {ciou_ab:.4f}\n") # IoU, DIoU, CIoU都为1

    # 场景2：两个框部分重叠，中心点有偏差，宽高比一致
    box_c = [0, 0, 10, 10]
    box_d = [2, 2, 12, 12] # 中心点有偏差，但都是正方形
    iou_cd = calculate_iou(box_c, box_d)
    ciou_cd = calculate_ciou(box_c, box_d)
    print(f"场景2：框C {box_c}, 框D {box_d} (部分重叠，中心点有偏差，宽高比一致)")
    print(f"IoU: {iou_cd:.4f}, CIoU: {ciou_cd:.4f}\n")

    # 场景3：两个框中心点重合，但宽高比不同 (DIoU可能高，CIoU会惩罚宽高比)
    box_e_fixed = [0, 0, 20, 10] # 中心点 (10, 5)
    box_f_fixed = [5, -5, 15, 15] # 宽10高20，中心点 (10, 5)
    iou_ef_fixed = calculate_iou(box_e_fixed, box_f_fixed)
    ciou_ef_fixed = calculate_ciou(box_e_fixed, box_f_fixed) # 修复了这里的bug，应该是box_f_fixed
    print(f"场景3：框E {box_e_fixed}, 框F {box_f_fixed} (中心点重合，宽高比不同)")
    print(f"IoU: {iou_ef_fixed:.4f}, CIoU: {ciou_ef_fixed:.4f}\n")

    # 场景4：两个框没有交集，离得很远 (IoU=0, DIoU和CIoU都是负值，CIoU的宽高比惩罚仍然存在)
    box_g = [0, 0, 5, 5]
    box_h = [100, 100, 105, 105]
    iou_gh = calculate_iou(box_g, box_h)
    ciou_gh = calculate_ciou(box_g, box_h)
    print(f"场景4：框G {box_g}, 框H {box_h} (无交集，离得很远)")
    print(f"IoU: {iou_gh:.4f}, CIoU: {ciou_gh:.4f}\n")
    
    # 场景5：一个框包含另一个框，且宽高比一致 (IoU>0, CIoU惩罚中心点距离，宽高比惩罚小)
    box_i = [0, 0, 20, 20]
    box_j = [5, 5, 15, 15]
    iou_ij = calculate_iou(box_i, box_j)
    ciou_ij = calculate_ciou(box_i, box_j)
    print(f"场景5：框I {box_i} (大框), 框J {box_j} (小框，被包含，宽高比一致)")
    print(f"IoU: {iou_ij:.4f}, CIoU: {ciou_ij:.4f}\n")

    # 场景6：一个框包含另一个框，但宽高比不同 (IoU>0, CIoU同时惩罚中心点距离和宽高比)
    box_k = [0, 0, 20, 20] # 正方形
    box_l = [5, 8, 15, 12] # 扁平长方形，中心点 (10, 10)
    iou_kl = calculate_iou(box_k, box_l)
    ciou_kl = calculate_ciou(box_k, box_l)
    print(f"场景6：框K {box_k} (大框), 框L {box_l} (小框，被包含，宽高比不同)")
    print(f"IoU: {iou_kl:.4f}, CIoU: {ciou_kl:.4f}\n")